import os

import matplotlib.pyplot as plt
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from skimage.metrics import structural_similarity as ssim
from sklearn.metrics import f1_score, jaccard_score
from torch.utils.data import DataLoader, Dataset
from torchvision import transforms

# 导入模型定义
from segmentation_train import ResNetSegmentation  # 确保你已经定义了模型并可以正确导入


# 定义测试数据集
class CustomTestDataset(Dataset):
    def __init__(self, input_dir, mask_dir, transform=None):
        self.input_dir = input_dir
        self.mask_dir = mask_dir
        self.image_names = os.listdir(input_dir)
        self.transform = transform

    def __len__(self):
        return len(self.image_names)

    def __getitem__(self, idx):
        img_name = self.image_names[idx]
        img_path = os.path.join(self.input_dir, img_name)
        mask_path = os.path.join(self.mask_dir, img_name)

        # 加载图像和掩码
        image = Image.open(img_path).convert('RGB')  # 确保图像为RGB模式
        mask = Image.open(mask_path).convert('L')  # 确保掩码为灰度模式

        # 应用图像变换
        if self.transform:
            image = self.transform(image)
            # 对掩码应用相同的尺寸变换，但不进行归一化
            mask_transform = transforms.Compose([
                transforms.Resize((256, 256)),
                transforms.ToTensor(),
            ])
            mask = mask_transform(mask)
        else:
            # 如果未提供transform，则手动进行转换
            image = transforms.ToTensor()(image)
            image = transforms.Resize((256, 256))(image)
            mask = transforms.ToTensor()(mask)
            mask = transforms.Resize((256, 256))(mask)

        return image, mask


# 评估指标函数
def evaluate_metrics(test_Y, pred_Y):
    mse = F.mse_loss(pred_Y, test_Y).item()
    test_Y_np = test_Y.squeeze(1).cpu().numpy()
    pred_Y_np = pred_Y.squeeze(1).cpu().numpy()

    # 计算SSIM，假设图像在[0,1]范围内
    ssim_score = np.mean([
        ssim(test_Y_np[i], pred_Y_np[i], data_range=1)
        for i in range(test_Y_np.shape[0])
    ])

    # 将预测和真实掩码二值化
    binary_test_Y = (test_Y_np > 0.5).astype(int)
    binary_pred_Y = (pred_Y_np > 0.5).astype(int)

    # 计算Dice系数和IoU
    dice_coefficient = np.mean([
        f1_score(binary_test_Y[i].flatten(), binary_pred_Y[i].flatten())
        for i in range(binary_test_Y.shape[0])
    ])
    iou = np.mean([
        jaccard_score(binary_test_Y[i].flatten(), binary_pred_Y[i].flatten())
        for i in range(binary_test_Y.shape[0])
    ])

    return mse, ssim_score, dice_coefficient, iou


# 显示图像函数
def display_images(test_X, test_Y, pred_Y):
    test_X_np = test_X[0].permute(1, 2, 0).cpu().numpy()
    test_Y_np = test_Y[0].squeeze(0).cpu().numpy()
    pred_Y_np = pred_Y[0].squeeze(0).cpu().numpy()

    # 反归一化以显示图像
    mean = np.array([0.485, 0.456, 0.406])
    std = np.array([0.229, 0.224, 0.225])
    test_X_display = np.clip(test_X_np * std + mean, 0, 1)

    plt.figure(figsize=(12, 4))
    plt.subplot(1, 3, 1)
    plt.imshow(test_X_display)
    plt.title("input image")
    plt.axis('off')

    plt.subplot(1, 3, 2)
    plt.imshow(test_Y_np, cmap='gray')
    plt.title("real mask")
    plt.axis('off')

    plt.subplot(1, 3, 3)
    plt.imshow(pred_Y_np, cmap='gray')
    plt.title("output mask")
    plt.axis('off')
    plt.show()


# 测试模型函数
def test_model(model, dataloader, device):
    model.eval()
    total_images = 0
    all_metrics = {'mse': [], 'ssim': [], 'dice': [], 'iou': []}

    with torch.no_grad():
        for i, (test_X, test_Y) in enumerate(dataloader):
            test_X, test_Y = test_X.to(device), test_Y.to(device)
            pred_Y = model(test_X)

            mse, ssim_score, dice_coefficient, iou = evaluate_metrics(test_Y, pred_Y)

            all_metrics['mse'].append(mse)
            all_metrics['ssim'].append(ssim_score)
            all_metrics['dice'].append(dice_coefficient)
            all_metrics['iou'].append(iou)

            total_images += test_X.size(0)
            print(f"样本 {i + 1}: MSE={mse:.4f}, SSIM={ssim_score:.4f}, Dice={dice_coefficient:.4f}, IoU={iou:.4f}")

            if i < 10:  # 只展示前10个批次的图像
                display_images(test_X, test_Y, pred_Y)

    avg_metrics = {metric: np.mean(values) for metric, values in all_metrics.items()}
    print("\n平均指标:")
    print(f"均方误差 (MSE): {avg_metrics['mse']:.4f}")
    print(f"结构相似性指数 (SSIM): {avg_metrics['ssim']:.4f}")
    print(f"Dice 系数: {avg_metrics['dice']:.4f}")
    print(f"交并比 (IoU): {avg_metrics['iou']:.4f}")
    print(f"测试的总图片数: {total_images}")


# 设置设备和加载模型
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = ResNetSegmentation(pretrained=False, num_classes=1)
model.load_state_dict(torch.load('../pt_file/resnet_segmentation_final.pth', map_location=device))
model.to(device)

# 定义图像变换
transform = transforms.Compose([
    transforms.Resize((256, 256)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
])

# 加载测试数据集
test_dataset = CustomTestDataset('../unet_img/imgs/test', '../unet_img/masks/test', transform=transform)
test_dataloader = DataLoader(test_dataset, batch_size=4, shuffle=False)

# 开始测试
test_model(model, test_dataloader, device)

# 平均指标:
# 均方误差 (MSE): 0.0059
# 结构相似性指数 (SSIM): 0.8925
# Dice 系数: 0.9757
# 交并比 (IoU): 0.9526
# 测试的总图片数: 400
